#ifndef YOLOV5_API_H
#define YOLOV5_API_H
#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>
namespace yolov5
{

    class Yolov5_api
    {

    public:
        Yolov5_api(const int gpu_id, const std::string &onnx_path, const std::string &trt_dir, const int max_batch = 1);

        Yolov5_api(const int gpu_id, const std::string &onnx_path, const std::string &trt_dir,
                   const std::vector<std::vector<std::vector<int>>> &anchor_size,
                   const int image_size,
                   const int max_batch = 1);

        ~Yolov5_api();

    public:
        std::vector<std::vector<float>> Execute(const cv::Mat &img);

        inline std::vector<std::string> &GetClassesName()
        {
            return class_name_;
        }

        inline void SetClassesName(const std::vector<std::string> &class_name)
        {
            class_name_ = class_name;
        }

    private:
        std::vector<std::vector<float>> ExecuteGPU(const cv::Mat &img);

        Yolov5_api();
        // std::vector<std::string> class_name_={"hat","head","Head gear","red hat","white peaked cap","chef hat","bar chef hat"};
        // std::vector<std::string> class_name_={"Regional Defects", "Black Long Warp", "White Weft", "Fuzzy Weft"};
        std::vector<std::string> class_name_ = {"Black Long Warp", "White Weft", "Fuzzy Weft"};

        class Impl;

        std::unique_ptr<Impl> impl_;
    };

} //namespace yolov5

#endif